import time
from dataset import get_tiny_imagenet_data
from utils import parse_option, build_resnet18_tiny_imagenet, build_swin_tiny_imagenet, build_vgg_tiny_imagenet,train_process, plot_results
from model.Efficientnet import create_efficientnet

def get_model(model_type):
    """根据模型类型构建并返回模型"""
    model = None
    try:
        if model_type == 'SwinTransformer':
            model = build_swin_tiny_imagenet()
        elif model_type == 'Vgg16':
            model = build_vgg_tiny_imagenet()
        elif model_type == 'ResNet18':
            model = build_resnet18_tiny_imagenet()
        elif model_type == 'EfficientNet':
            model = create_efficientnet()
        else:
            # 明确抛出异常，说明支持的模型类型
            raise ValueError(
                f"Unsupported model_type: '{model_type}'. "
                "Expected 'SwinTransformer' or 'Vgg16' or 'ResNet18'"
            )
    except Exception as e:
        # 捕获可能的子函数错误（如build_xxx失败）
        raise RuntimeError(f"Failed to build model '{model_type}': {str(e)}")

    return model


if __name__ == "__main__":
    args = parse_option()
    model = get_model(args.model_type)

    #用于实验的tiny ImageNet数据集
    train_loader, val_loader, num_classes = get_tiny_imagenet_data(batch_size=args.batch_size,
                                                                   num_workers=args.num_workers)

    history = train_process(model, train_loader, val_loader, epochs=args.epochs, gpu=args.gpu, lr=args.lr, weight_decay=args.weight_decay)
    # 可视化
    plot_results(history, args.model_type)

